import sys
import os
import io
import re
import pandas as pd
import chardet
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QVBoxLayout, QHBoxLayout, QWidget,
    QTreeView, QFileSystemModel, QPushButton,
    QListWidget, QListWidgetItem, QFileDialog, QTreeWidget, QTreeWidgetItem, QMenu
)
from PySide6.QtCore import Qt, QEvent,QAbstractTableModel
from PySide6.QtCharts import QChart, QChartView, QLineSeries, QValueAxis
from PySide6.QtGui import QPainter, QPen, QColor
import random


class DataFrameModel(QAbstractTableModel):
    def __init__(self, df):
        super().__init__()
        self.df = df

    def rowCount(self, parent=None):
        return len(self.df)

    def columnCount(self, parent=None):
        return len(self.df.columns)

    def data(self, index, role=Qt.DisplayRole):
        if role == Qt.DisplayRole:
            return str(self.df.iloc[index.row(), index.column()])
        return None

    def headerData(self, section, orientation, role=Qt.DisplayRole):
        if role == Qt.DisplayRole:
            if orientation == Qt.Horizontal:
                return str(self.df.columns[section])
            elif orientation == Qt.Vertical:
                return str(self.df.index[section])
        return None


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.series_list = []  # 存储图表系列
        self.initUI()

    def initUI(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        # 水平布局
        hbox = QHBoxLayout(central_widget)

        # 文件选择按钮
        self.select_file_button = QPushButton("选择文件")
        self.select_file_button.clicked.connect(self.select_file)
        hbox.addWidget(self.select_file_button)

        # 左侧树视图
        self.tree_widget = QTreeWidget()
        self.tree_widget.setHeaderLabels(["井名及表头"])
        hbox.addWidget(self.tree_widget, 1)

        # 右侧图表区域
        self.chart_widget = QWidget()
        chart_layout = QVBoxLayout(self.chart_widget)
        self.chart_view = QChartView()
        self.chart_view.setDragMode(QChartView.ScrollHandDrag)  # 允许拖动图表
        chart_layout.addWidget(self.chart_view)
        hbox.addWidget(self.chart_widget, 3)

        # 连接树的上下文菜单请求信号
        self.tree_widget.customContextMenuRequested.connect(self.show_context_menu)
        self.tree_widget.itemChanged.connect(self.update_chart_on_check)  # 连接勾选变化信号

        self.setWindowTitle("多格式文件读取与树结构展示")
        self.resize(1200, 800)

    def select_file(self):
        file_dialog = QFileDialog()
        file_path, _ = file_dialog.getOpenFileName(self, "选择文件", "", "文本文件 (*.txt *.csv)")
        if file_path:
            try:
                print(f"当前文件路径: {file_path}")
                file_name = os.path.basename(file_path)
                well_name = self.extract_well_name(file_name)
                print(f"当前文件所属井名: {well_name}")
                self.df = self.read_file(file_path)

                # 验证数据是否读取成功
                if self.df is not None and not self.df.empty:
                    print(f"成功读取数据，列名: {self.df.columns.tolist()}")
                    self.populate_tree(well_name)
                else:
                    print("警告：读取的数据为空")
            except Exception as e:
                print(f"读取文件时出错: {e}")

    def read_file(self, file_path):
        try:
            # 检测文件编码
            with open(file_path, 'rb') as f:
                raw_data = f.read()
                result = chardet.detect(raw_data)
                encoding = result['encoding'] or 'gbk'  # 防止编码检测为None
                if encoding == 'ISO-8859-1':
                    encoding = 'gbk'
                print(f"文件编码: {encoding}")

            # 优先处理中文编码
            if encoding in ['ISO-8859-1', 'ascii']:
                encoding = 'gbk'

            # 尝试三种格式读取
            df = self.try_three_formats(file_path, encoding)
            if not df.empty:
                return df

            raise ValueError("未匹配到任何文件格式")

        except Exception as e:
            print(f"读取文件 {file_path} 时出错: {e}")
            return pd.DataFrame()

    def try_three_formats(self, file_path, encoding):
        # 格式1: 井深 钻时 全烃（简单表格）
        try:
            df = pd.read_csv(file_path, sep=r'\s+', header=0, encoding=encoding)
            if not df.empty:
                return df
        except:
            pass

        # 格式2: 井号 深度 岩性简述（简单表格）
        try:
            df = pd.read_csv(file_path, sep=r'\s+', header=0, encoding=encoding)
            if not df.empty:
                return df
        except:
            pass

        # 格式3: 包含 ~ASCII DATA SECTION（复杂格式）
        with open(file_path, 'r', encoding=encoding) as f:
            lines = f.readlines()
            data_start = next((i for i, line in enumerate(lines) if line.strip().startswith('~ASCII DATA SECTION')),
                              None)
            if data_start is None:
                return pd.DataFrame()

            # 提取表头（假设表头在 DATA SECTION 前2行）
            header_line = lines[data_start - 2].strip().lstrip('#').strip()
            headers = header_line.split() if header_line else []

            # 提取数据行
            data_lines = [line for line in lines[data_start + 1:] if line.strip()]
            if not data_lines:
                return pd.DataFrame()

            # 解析数据
            data_str = '\n'.join(data_lines)
            df = pd.read_csv(io.StringIO(data_str), sep=r'\s+', header=None, names=headers)
            return df

    def extract_well_name(self, file_name):
        # # 优化井名提取逻辑
        # pattern = r'(\d+-\d+井)'  # 匹配如 3-6井
        # match = re.search(pattern, file_name)
        # return match.group(1) if match else file_name.split('.')[0]  # 兜底方案
        pattern = r'(\d+-\d+[A-Za-z]+)'
        match = re.search(pattern, file_name)
        if match:
            return match.group(1) + '井'
        index = file_name.find('井')
        if index != -1:
            return file_name[:index + 1]
        return file_name.split('.')[0]

    # def populate_tree(self, well_name):
    #     self.tree_widget.clear()
    #     if well_name:
    #         parent_item = QTreeWidgetItem(self.tree_widget, [well_name])
    #         curve_collection = QTreeWidgetItem(parent_item, ["曲线集合"])
    #
    #         if hasattr(self, 'df') and not self.df.empty:
    #             for col in self.df.columns:
    #                 child_item = QTreeWidgetItem(curve_collection, [col])
    #                 child_item.setCheckState(0, Qt.Unchecked)  # 设置子节点可勾选
    #         self.tree_widget.expandAll()

    def show_context_menu(self, pos):
        item = self.tree_widget.itemAt(pos)
        if item and item.text(0) == "曲线集合":
            menu = QMenu()
            import_action = menu.addAction("导入表信息")
            import_action.triggered.connect(self.import_table_info)
            menu.exec_(self.tree_widget.mapToGlobal(pos))

    def import_table_info(self):
        if hasattr(self, 'df') and not self.df.empty:
            well_name = self.extract_well_name('') or "未知井"  # 确保井名存在
            self.populate_tree(well_name)
            self.update_chart([col for col in self.df.columns])  # 导入时显示所有列

    def update_chart_on_check(self, item, column):
        if item.parent() and item.parent().text(0) == "曲线集合":
            checked_headers = self.collect_checked_headers()
            self.update_chart(checked_headers)

    def collect_checked_headers(self):
        checked = []
        root_item = self.tree_widget.invisibleRootItem()
        for i in range(root_item.childCount()):  # 使用 childCount() 和 child(i) 遍历子节点
            well_item = root_item.child(i)  # 井名节点
            for j in range(well_item.childCount()):
                curve_collection = well_item.child(j)  # 曲线集合节点
                if curve_collection.text(0) == "曲线集合":
                    for k in range(curve_collection.childCount()):
                        col_item = curve_collection.child(k)
                        if col_item.checkState(0) == Qt.Checked:
                            checked.append(col_item.text(0))
        return checked

    def populate_tree(self, well_name):
        self.tree_widget.clear()
        if well_name:
            well_item = QTreeWidgetItem(self.tree_widget, [well_name])  # 明确井名节点
            curve_collection = QTreeWidgetItem(well_item, ["曲线集合"])  # 曲线集合作为井名的子节点
            if hasattr(self, 'df') and not self.df.empty:
                for col in self.df.columns:
                    child_item = QTreeWidgetItem(curve_collection, [col])
                    child_item.setCheckState(0, Qt.Unchecked)
            self.tree_widget.expandAll()

    def update_chart(self, headers):
        chart = QChart()
        chart.setTitle("Data from Table")

        # 设置坐标轴（关键：避免数据范围未显示）
        x_axis = QValueAxis()
        y_axis = QValueAxis()
        chart.addAxis(x_axis, Qt.AlignBottom)
        chart.addAxis(y_axis, Qt.AlignLeft)

        # 清除之前的系列
        for series in self.series_list:
            chart.removeSeries(series)
        self.series_list = []

        for col in headers:
            if col in self.df.columns:
                series = QLineSeries()
                for idx, value in self.df[col].items():  # 使用items()获取索引和值
                    try:
                        series.append(idx, float(value))  # idx作为x轴，数据作为y轴
                    except ValueError:
                        print(f"警告：列 {col} 包含非数值数据: {value}")
                        continue
                if series.count() > 0:  # 确保有数据点
                    series.setName(col)
                    series.setPen(self.get_random_pen())
                    chart.addSeries(series)
                    self.series_list.append(series)
                    series.attachAxis(x_axis)  # 关联坐标轴
                    series.attachAxis(y_axis)

        self.chart_view.setChart(chart)
        self.chart_view.setRenderHint(QPainter.Antialiasing)
        self.chart_view.setInteractive(True)
        self.chart_view.chart().setToolTip("Hover over a point to see data")
        self.chart_view.installEventFilter(self)

    def eventFilter(self, source, event):
        if source == self.chart_view and event.type() == QEvent.MouseMove:
            pos = event.pos()
            chart = self.chart_view.chart()
            for series in chart.series():
                for i in range(series.count()):
                    point = series.at(i)
                    screen_pos = self.chart_view.mapFromParent(chart.mapToPosition(point))  # 简化坐标转换
                    if (pos - screen_pos).manhattanLength() < 10:
                        chart.setToolTip(f"{series.name()}: x={point.x()}, y={point.y()}")
                        return True
            chart.setToolTip("")
        return super().eventFilter(source, event)

    def get_random_pen(self):
        r = random.randint(0, 255)
        g = random.randint(0, 255)
        b = random.randint(0, 255)
        return QPen(QColor(r, g, b))


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())